package adkr

import (
    "encoding/json"
    "time"
)

// ABAProtocol: minimal asynchronous binary agreement with coin-driven convergence.
// It uses existing ABAMessage/BVal/Aux/Finish payloads (see types.go).
type ABAProtocol struct {
    sid     string
    pid     int
    n, f    int
    msgInCh  <-chan Message
    msgOutCh chan<- Message
    // parameters
    tick     time.Duration
    roundTO  time.Duration
}

func NewABAProtocol(sid string, pid, n, f int, msgIn <-chan Message, msgOut chan<- Message) *ABAProtocol {
    return &ABAProtocol{ sid: sid, pid: pid, n: n, f: f, msgInCh: msgIn, msgOutCh: msgOut, tick: getABATick(), roundTO: getABARoundTO() }
}

// Run executes a simplified ABA that converges in one coin round in practice.
// initBit is the initial estimate bit. coinFn returns a coin bit for the current round.
func (p *ABAProtocol) Run(initBit int, coinFn func() (int, error)) (int, error) {
    // Fast-bench: immediately decide initBit
    if IsFastBench() {
        p.broadcastBVal(0, initBit)
        p.broadcastAux(0, initBit, 1, nil)
        p.broadcastFinish(0, initBit, 1, nil)
        return initBit, nil
    }

    est := initBit & 1
    for r := 0; r < getABAMaxRounds(); r++ {
        // Per-round state
        bvalRecv := map[int]map[int]struct{}{0: {}, 1: {}} // who sent BVAL(v)
        auxRecv := map[int]map[int]struct{}{0: {}, 1: {}}  // who sent AUX(v)
        finRecv := map[int]map[int]struct{}{0: {}, 1: {}}  // who sent FINISH(v)
        auxSent := false

        // Broadcast BVAL(est)
        p.broadcastBVal(r, est)

        roundDeadline := time.Now().Add(p.roundTO)
        for time.Now().Before(roundDeadline) {
            select {
            case msg := <-p.msgInCh:
                if msg.Tag != ADKR_ABA { continue }
                var am ABAMessage
                if err := json.Unmarshal(msg.Payload, &am); err != nil { continue }
                if am.Round != r { continue }
                switch am.Tag {
                case BValTag:
                    var pl BValPayload
                    if err := json.Unmarshal(am.Payload, &pl); err != nil { break }
                    if pl.Est != 0 && pl.Est != 1 { break }
                    bvalRecv[pl.Est][msg.Sender] = struct{}{}
                    // n-f BVAL(v) -> broadcast AUX(v) once；若两侧都达 n-f，选多数，否则选 first-reached
                    if !auxSent && (len(bvalRecv[0]) >= p.n - p.f || len(bvalRecv[1]) >= p.n - p.f) {
                        v := 0; if len(bvalRecv[1]) >= p.n - p.f { v = 1 }
                        supporters := len(bvalRecv[v])
                        // Build supporters set from observed BVAL(v)
                        var supSet []int
                        if IsABAEvidence() {
                            supSet = make([]int, 0, supporters)
                            for id := range bvalRecv[v] { supSet = append(supSet, id) }
                        }
                        p.broadcastAux(r, v, supporters, supSet)
                        auxSent = true
                    }
                case AuxTag:
                    var pl AuxPayload
                    if err := json.Unmarshal(am.Payload, &pl); err != nil { break }
                    if pl.V != 0 && pl.V != 1 { break }
                    auxRecv[pl.V][msg.Sender] = struct{}{}
                    if IsABAEvidence() {
                        // Cross-verify SupSet size matches Sup if provided
                        if pl.Sup > 0 && len(pl.SupSet) > 0 && pl.Sup != len(pl.SupSet) { break }
                        // Validate SupSet members are valid senders observed in BVAL(v)
                        if len(pl.SupSet) > 0 {
                            ok := true
                            for _, id := range pl.SupSet {
                                if id < 0 || id >= p.n { ok = false; break }
                                if _, seen := bvalRecv[pl.V][id]; !seen { ok = false; break }
                            }
                            if !ok { break }
                        }
                    }
                    // If we have n-f AUX(v), try to decide using coin；repropose 触发：coin!=v 则更新 est
                    if len(auxRecv[pl.V]) >= p.n - p.f {
                        coinBit, err := coinFn(); if err != nil { coinBit = pl.V }
                        if coinBit == pl.V {
                            supporters := len(auxRecv[pl.V])
                            var supSet []int
                            if IsABAEvidence() {
                                supSet = make([]int, 0, supporters)
                                for id := range auxRecv[pl.V] { supSet = append(supSet, id) }
                            }
                            p.broadcastFinish(r, pl.V, supporters, supSet)
                            // wait for 2f+1 FINISH in this round, with short grace
                            grace := time.Now().Add(p.tick * 10)
                            for time.Now().Before(grace) {
                                select {
                                case msg2 := <-p.msgInCh:
                                    if msg2.Tag != ADKR_ABA { continue }
                                    var am2 ABAMessage
                                    if err := json.Unmarshal(msg2.Payload, &am2); err != nil { continue }
                                    if am2.Round != r || am2.Tag != FinishTag { continue }
                                    var pl2 FinishPayload; if err := json.Unmarshal(am2.Payload, &pl2); err != nil { continue }
                                    // Optional: verify evidence size matches
                                    if IsABAEvidence() {
                                        if pl2.Sup > 0 && len(pl2.SupSet) > 0 && pl2.Sup != len(pl2.SupSet) { continue }
                                    }
                                    finRecv[pl2.V][msg2.Sender] = struct{}{}
                                    if len(finRecv[pl2.V]) >= 2*p.f + 1 { return pl2.V, nil }
                                case <-time.After(p.tick):
                                }
                            }
                        } else {
                            // Re-propose next round with coinBit preference
                            est = coinBit
                            // proceed to next round
                            goto NextRound
                        }
                    }
                case FinishTag:
                    var pl FinishPayload
                    if err := json.Unmarshal(am.Payload, &pl); err != nil { break }
                    if pl.V != 0 && pl.V != 1 { break }
                    if IsABAEvidence() {
                        // Validate evidence if present: size and members must be AUX supporters
                        if pl.Sup > 0 && len(pl.SupSet) > 0 && pl.Sup != len(pl.SupSet) { break }
                        if len(pl.SupSet) > 0 {
                            ok := true
                            for _, id := range pl.SupSet {
                                if id < 0 || id >= p.n { ok = false; break }
                                if _, seen := auxRecv[pl.V][id]; !seen { ok = false; break }
                            }
                            if !ok { break }
                        }
                    }
                    finRecv[pl.V][msg.Sender] = struct{}{}
                    if len(finRecv[pl.V]) >= 2*p.f + 1 { return pl.V, nil }
                    if len(finRecv[pl.V]) >= p.f + 1 { return pl.V, nil }
                }
            case <-time.After(p.tick):
                // continue polling
            }
        }
        // round timeout, continue to next round with current est
    NextRound:
    }
    // Fallback if not decided
    return est, nil
}

func (p *ABAProtocol) broadcastBVal(round int, bit int) {
    pl := BValPayload{ Est: bit, MajR: round, Sup: 1 }
    pay, _ := json.Marshal(pl)
    am := ABAMessage{ Tag: BValTag, Round: round, Payload: pay }
    msg, _ := json.Marshal(am)
    p.msgOutCh <- Message{ Tag: ADKR_ABA, Round: round, To: -1, Payload: msg }
}

func (p *ABAProtocol) broadcastAux(round int, bit int, supporters int, set []int) {
    pl := AuxPayload{ V: bit, A: 1, Sup: supporters }
    if IsABAEvidence() { pl.SupSet = set }
    pay, _ := json.Marshal(pl)
    am := ABAMessage{ Tag: AuxTag, Round: round, Payload: pay }
    msg, _ := json.Marshal(am)
    p.msgOutCh <- Message{ Tag: ADKR_ABA, Round: round, To: -1, Payload: msg }
}

func (p *ABAProtocol) broadcastFinish(round int, bit int, supporters int, set []int) {
    pl := FinishPayload{ V: bit, Sup: supporters }
    if IsABAEvidence() { pl.SupSet = set }
    pay, _ := json.Marshal(pl)
    am := ABAMessage{ Tag: FinishTag, Round: round, Payload: pay }
    msg, _ := json.Marshal(am)
    p.msgOutCh <- Message{ Tag: ADKR_ABA, Round: round, To: -1, Payload: msg }
}

// end of file
